#include <vector>
#include <torch/extension.h>


// C++ API of CUDA forward & backward declarations
std::vector<torch::Tensor> costh_cuda_fwd(const torch::Tensor, const torch::Tensor);
std::vector<torch::Tensor> costh_cuda_bkwd(const torch::Tensor, const torch::Tensor, 
                                           const torch::Tensor, const torch::Tensor, 
                                           const torch::Tensor);


// C++ high level interface declarations & macro definition
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

// When testing, variable `x_norm` shape: (batch_size, 768);
// variable `w` shape: (768, 1211)
std::vector<torch::Tensor> costh_fwd(torch::Tensor x_norm, torch::Tensor w)
{
    CHECK_INPUT(x_norm);
    CHECK_INPUT(w);

    return costh_cuda_fwd(x_norm, w);
}

// When testing, variable `grad_mm` shape: (batch_size, 1211)
// variable `x_norm` shape: (batch_size, 768)
// variable `w` shape: (768, 1211)
// variable `costh_out` shape: (768, 1211)
// variable `norm_vals` shape: (1211, )

std::vector<torch::Tensor> costh_bkwd(torch::Tensor grad_mm, torch::Tensor x_norm, 
                                      torch::Tensor w, torch::Tensor costh_out, 
                                      torch::Tensor norm_vals)
{
    CHECK_INPUT(grad_mm);
    CHECK_INPUT(x_norm);
    CHECK_INPUT(w);
    CHECK_INPUT(costh_out);
    CHECK_INPUT(norm_vals);

    return costh_cuda_bkwd(grad_mm, x_norm, w, costh_out, norm_vals);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("forward", &costh_fwd, "the forward of cosine similarity computation in Multitask-W2V (CUDA)\n"
                                  "There're two input Tensor parameters of this forward function: \n"
                                  "(1) x_norm: (batch_size, 768); (2) w: (768, 1211).\n"
                                  "On the other hand, there're three return values of this forward function: \n"
                                  "(1) mm_out: (batch_size, 1211); (2) costh_out: (768, 1211); (3) norm_vals: (1211, ).");
    m.def("backward", &costh_bkwd, "the backward of cosine similarity computation in Multitask-W2V (CUDA)\n"
                                   "There're five input Tensor parameters of this backward function: \n"
                                   "(1) grad_mm: (batch_size, 1211); (2) x_norm: (batch_size, 768); (3) w: (768, 1211); \n"
                                   "(4) costh_out: (768, 1211); (5) norm_vals: (1211, ).\n"
                                   "On the other hand, there're two return values of backward function: \n"
                                   "(1) grad_x_norm: (batch_size, 768); (2) grad_w: (768, 1211).");
}
